import akshare as ak
import pandas as pd
import time
import random
import os
import logging
from tqdm import tqdm
from datetime import datetime

# 日志配置
logging.basicConfig(
    filename='data/akshare_errors.log',
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

output_dir = "data/akshare_parquet_data"
os.makedirs(output_dir, exist_ok=True)

def get_existing_stock_codes(directory):
    """
    获取指定目录下已有的股票代码（从文件名中提取）。
    """
    existing_codes = set()
    if os.path.exists(directory):
        for filename in os.listdir(directory):
            if filename.endswith(".parquet"):
                code = filename.replace(".parquet", "")
                existing_codes.add(code)
    print(f"在 {directory} 中发现 {len(existing_codes)} 个已存在的股票文件。")
    return existing_codes

def get_all_stock_codes():
    """
    获取所有 A 股股票代码。
    """
    try:
        df = ak.stock_info_a_code_name()
        codes = df["code"].tolist()
        print(f"共获取 {len(codes)} 个股票代码")
        return codes
    except Exception as e:
        logging.error(f"获取股票代码失败：{e}")
        return []

def download_stock_with_retry(code, max_retries=5):
    """
    下载单个股票的历史数据，支持重试机制。
    """
    if code.startswith("6"):
        stock_symbol = f"sh{code}"
    else:
        stock_symbol = f"sz{code}"

    # 获取今天的日期并格式化为字符串（例如："20250612"）
    today = datetime.now().strftime("%Y%m%d")
    
    # 设置下载的日期范围（开始日期和结束日期都为今天）
    start_date = today
    end_date = today
    
    print(f"start_date: {start_date}")  # 输出示例：start_date: 20250612
    print(f"end_date: {end_date}")      # 输出示例：end_date: 20250612

    for attempt in range(1, max_retries + 1):
        try:
            print(f"正在下载 {stock_symbol} - 尝试 {attempt}")
            df = ak.stock_zh_a_daily(symbol=stock_symbol, start_date=start_date, end_date=end_date)

            if df.empty:
                logging.info(f"{code} 数据为空，跳过")
                return None

            # 数据清洗和格式化
            df["capital"] = (df["outstanding_share"] / 1e8).round(2) # 股本单位亿
            df["volume"] = (df["volume"] / 1e8).round(2)             # 成交量单位亿
            df["amount"] = (df["amount"] / 1e8).round(2)             # 成交额单位亿

            result_df = df[[
                "date", "open", "high", "low", "close",
                "volume", "amount", "capital", "turnover"
            ]]
            return result_df

        except Exception as e:
            logging.warning(f"{code} 第 {attempt} 次下载失败: {e}")
            time.sleep(3 + attempt * 2)  # 指数退避机制，等待更长时间再重试

    logging.error(f"{code} 连续下载失败 {max_retries} 次，已放弃。")
    return None

def main():
    """
    主函数，协调股票数据的下载。
    """
    all_stock_codes = get_all_stock_codes()
    if not all_stock_codes:
        print("股票列表为空，程序退出。")
        return

    existing_codes = get_existing_stock_codes(output_dir)

    # 过滤掉已存在的股票代码
    codes_to_download = [code for code in all_stock_codes if code not in existing_codes]
    print(f"将下载 {len(codes_to_download)} 个新的股票数据。")

    if not codes_to_download:
        print("所有股票数据已下载，无需更新。")
        return

    # 使用 tqdm 显示下载进度
    for code in tqdm(codes_to_download, desc="下载股票数据"):
        df = download_stock_with_retry(code)
        if df is not None:
            # 确保日期列为 datetime 类型以便后续处理
            df["date"] = pd.to_datetime(df["date"])
            output_path = os.path.join(output_dir, f"{code}.parquet")
            try:
                df.to_parquet(output_path, engine="fastparquet", index=False)
                print(f"成功保存 {code}.parquet")
            except Exception as e:
                logging.error(f"保存 {code}.parquet 失败: {e}")
        
        # 每次请求后延迟 1 到 3 秒，防止触发反爬机制
        time.sleep(random.uniform(1, 3))

if __name__ == "__main__":
    main()